/**
**! Andrew Windsor | January 2014
**!
**! This program is a simple steganograhpy terminal-based application. It can take a txt file
**! and a bmp image file and encode the txt file into the image, providing a rudimentary form
**! hiding messages or other text. The program can then extract and reconstruct the original
**! text from an altered image. It was created as a proof of concept in studying basic computer
**! cryptography.
**!
**! In order to achieve the encoding the program uses a simple version of bitmasking the
**! least significant bit (LSB) within the image's bitstream.
**!
**! Uses BMP image files and TXT files. Other formats are not expected to work properly or at all.
**! 
**! This program comes with no warranty or gurantee of proper use. Use at your own risk.
**! Feel free to edit, change, improve, or study this program in any way that you'd like. 
*/

#include <iostream>
#include <fstream>
#include <bitset>
#include <string>
using namespace std;

void stegoDraw();

int main()
{
	stegoDraw(); //A bit of ASCII art for fun
	std::cout << "(STEGO)NOGRAPHY" << std::endl;
	std::cout << "" << std::endl;
	
	int choice = 0;
	std::cout << "Select operation to perform:" << std::endl;
	std::cout << "(1) Embed file" << std::endl;
	std::cout << "(2) Extract file" << std::endl;
	std::cout << "User choice: ";
	std::cin >> choice;
	
	switch (choice)
	{
		case 1: //function for embedding a file into an image
		{
			string filePathText;
			string imagePathText;
			streampos fileBufferSize;
			streampos imageBufferSize;
			char* embedStream;
			char* hostStream;
			
			//receive file paths for use
			std::cout << "Usage: ImagePath FilePath" << std::endl;
			std::cin >> imagePathText >> filePathText;
			
			//define and initialize the file streams
			ifstream imagePath(imagePathText.data(), ios::binary | ios::ate);
			ifstream filePath(filePathText.data(), ios::binary | ios::ate);
			
			//perform steganography operations on the two files
			
			//reading input file for embedding
			if (filePath.is_open())
			{
				fileBufferSize = filePath.tellg();
				embedStream = new char[fileBufferSize];
				filePath.seekg(0, ios::beg);
				filePath.read(embedStream, fileBufferSize);
				filePath.close();
			}
			else
			{
				std::cout << "Error, unable to open file." << std::endl;
			}
			//file input into image using LSB methodology
			if (imagePath.is_open())
			{
				//set up image mirror to manipulate bits
				imageBufferSize = imagePath.tellg();
				hostStream = new char[imageBufferSize];
				imagePath.seekg(0, ios::beg);
				imagePath.read(hostStream, imageBufferSize);
				imagePath.close();
				
				int imagePathCount = 60;//pad enough buffer in case of BMP header information
				std::bitset<8> fileBit;
				std::bitset<8> endFileBit;				
				for (int count = 0; count < fileBufferSize; count++)
				{
					fileBit = embedStream[count];
					for (int imageCount = 0; imageCount < 8; imageCount++)
					{
						hostStream[imagePathCount + imageCount] = (hostStream[imagePathCount + imageCount] & ~1) | fileBit[imageCount];	
					}
					imagePathCount += 8;
				}
				endFileBit = 0xFF;//2s compliment bit representation of -1 (fake end of file)
				for (int imageCount = 0; imageCount < 8; imageCount++)//add a fake end of file marker for the retrieval program 
				{
						hostStream[imagePathCount + imageCount] = (hostStream[imagePathCount + imageCount] & ~1) | endFileBit[imageCount];	
				}
				//reopen for output
				imagePathText.append(".output");//ensures program doesn't write over original file
				ofstream imageCopy(imagePathText.data(), ios::binary);
				imageCopy.write(hostStream, imageBufferSize);
				imageCopy.close();
			}
			else
			{
				std::cout << "Error, unable to write to file." << std::endl;
			}
			break;
		}
		
		case 2: //function for extracting a file from an embedded image
		{
			string filePathText;
			string imagePathText;
			streampos fileBufferSize;
			streampos imageBufferSize;
			char* embedStream;
			char* hostStream;
			
			//get input
			std::cout << "Usage: ImagePath NewFileName.txt" << std::endl;
			std::cin >> imagePathText >> filePathText;
			
			//define and initialize the image file stream
			ifstream imagePath(imagePathText.data(), ios::binary | ios::ate);
			ofstream filePath(filePathText.data(), ios::binary);
			
			//buffer the image for file extraction
			if (imagePath.is_open())
			{
				imageBufferSize = imagePath.tellg();
				hostStream = new char[imageBufferSize];
				imagePath.seekg(0, ios::beg);
				imagePath.read(hostStream, imageBufferSize);
				imagePath.close();
				
				int imagePathCount = 60; //match the starting position of the encoding function
				std::bitset<8> fileBit;
				
				//reconstruct the plain text file
				if (filePath.is_open())
				{
					while (fileBit != 0xFF)
					{
						for (int count = 0; count < 8; count++)
						{
							fileBit[count] = hostStream[imagePathCount + count] & 1;
						}
						unsigned long bitStream = fileBit.to_ulong(); //conversion of bitstream to readable type
						char convertedChar = static_cast<char>(bitStream); //conversion of long type to binary char
						filePath.put(convertedChar);
						imagePathCount += 8;
					}
					filePath.close();
				}
				else
				{
					std::cout << "Error, unable to write to file." << std::endl;
				}
			}
			else
			{
				std::cout << "Error, unable to open file." << std::endl;
			}
			
			break;
		}
		
		default: std::cout << "Invalid choice." << std::endl;
	}
	//program closes
	std::cout << "Program has closed successfully." << std::endl;
	return(0);
}

void stegoDraw()
{
	std::cout << "                           / `.   .\' \\" << std::endl;
	std::cout << "                   .---.  <    > <    >  .---." << std::endl;
	std::cout << "                   |    \\  \\ - ~ ~ - /  /    |" << std::endl;
	std::cout << "                    ~-..-~             ~-..-~" << std::endl;
	std::cout << "                \\~~~\\.\'                    `./~~~/" << std::endl;
	std::cout << "      .-~~^-.    \\__/                        \\__/" << std::endl;
	std::cout << "    .\'  O    \\     /               /       \\  \\" << std::endl;
	std::cout << "   (_____,    `._.\'               |         }  \\/~~~/" << std::endl;
	std::cout << "    `----.          /       }     |        /    \\__/" << std::endl;
	std::cout << "          `-.      |       /      |       /      `. ,~~|" << std::endl;
	std::cout << "              ~-.__|      /_ - ~ ^|      /- _      `..-\'   f: f:" << std::endl;
	std::cout << "                   |     /        |     /     ~-.     `-. _||_||_" << std::endl;
	std::cout << "                   |_____|        |_____|         ~ - . _ _ _ _ _>" << std::endl;
}